{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# An overview of QuantTensor and QuantConv2d\n",
    "\n",
    "In this initial tutorial, we take a first look at `QuantTensor`, a basic data structure in Brevitas, and at `QuantConv2d`, a typical quantized layer. `QuantConv2d` is an instance of a `QuantWeightBiasInputOutputLayer` (typically imported as `QuantWBIOL`), meaning that it supports quantization of its weight, bias, input and output. Other instances of `QuantWBIOL` are `QuantLinear`, `QuantConv1d`, `QuantConvTranspose1d` and `QuantConvTranspose2d`, and they all follow the same principles.\n",
    "\n",
    "If we take a look at the `__init__` method of `QuantConv2d`, we notice a few things:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/.local/lib/python3.7/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "```python\n",
       "    def __init__(\n",
       "            self,\n",
       "            in_channels: int,\n",
       "            out_channels: int,\n",
       "            kernel_size: Union[int, Tuple[int, int]],\n",
       "            stride: Union[int, Tuple[int, int]] = 1,\n",
       "            padding: Union[int, Tuple[int, int]] = 0,\n",
       "            dilation: Union[int, Tuple[int, int]] = 1,\n",
       "            groups: int = 1,\n",
       "            bias: bool = True,\n",
       "            padding_type: str = 'standard',\n",
       "            weight_quant: Optional[WeightQuantType] = Int8WeightPerTensorFloat,\n",
       "            bias_quant: Optional[BiasQuantType] = None,\n",
       "            input_quant: Optional[ActQuantType] = None,\n",
       "            output_quant: Optional[ActQuantType] = None,\n",
       "            return_quant_tensor: bool = False,\n",
       "            **kwargs) -> None:\n",
       "        Conv2d.__init__(\n",
       "            self,\n",
       "            in_channels=in_channels,\n",
       "            out_channels=out_channels,\n",
       "            kernel_size=kernel_size,\n",
       "            stride=stride,\n",
       "            padding=padding,\n",
       "            dilation=dilation,\n",
       "            groups=groups,\n",
       "            bias=bias)\n",
       "        QuantWBIOL.__init__(\n",
       "            self,\n",
       "            weight_quant=weight_quant,\n",
       "            bias_quant=bias_quant,\n",
       "            input_quant=input_quant,\n",
       "            output_quant=output_quant,\n",
       "            return_quant_tensor=return_quant_tensor,\n",
       "            **kwargs)\n",
       "        assert self.padding_mode == 'zeros'\n",
       "        assert not (padding_type == 'same' and padding != 0)\n",
       "        self.padding_type = padding_type\n",
       "\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import inspect\n",
    "from brevitas.nn import QuantConv2d\n",
    "from brevitas.nn import QuantIdentity\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "def pretty_print_source(source):\n",
    "    display(Markdown('```python\\n' + source + '\\n```'))\n",
    "    \n",
    "source = inspect.getsource(QuantConv2d.__init__)  \n",
    "pretty_print_source(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantConv2d` is an instance of both `Conv2d` and `QuantWBIOL`. Its initialization method exposes the usual arguments of a `Conv2d`, as well as: an extra flag to support *same padding*; *four* different arguments to set a quantizer for - respectively - *weight*, *bias*, *input*, and *output*; a `return_quant_tensor` boolean flag; the `**kwargs` placeholder to intercept additional arbitrary keyword arguments.  \n",
    "In this tutorial we will focus on how to set the four quantizer arguments and the return flags; arbitrary kwargs will be explained in a separate tutorial dedicated to defining and overriding quantizers.\n",
    "\n",
    "By default `weight_quant=Int8WeightPerTensorFloat`, while `bias_quant`, `input_quant` and `output_quant` are set to `None`. That means that by default weights are quantized to *8-bit signed integer with a per-tensor floating-point scale factor* (a very common type of quantization adopted by e.g. the ONNX standard opset), while quantization of bias, input, and output are disabled. We can easily verify all of this at runtime on an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "default_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Is weight quant enabled: True\n",
      "Is bias quant enabled: False\n",
      "Is input quant enabled: False\n",
      "Is output quant enabled: False\n"
     ]
    }
   ],
   "source": [
    "print(f'Is weight quant enabled: {default_quant_conv.is_weight_quant_enabled}')\n",
    "print(f'Is bias quant enabled: {default_quant_conv.is_bias_quant_enabled}')\n",
    "print(f'Is input quant enabled: {default_quant_conv.is_input_quant_enabled}')\n",
    "print(f'Is output quant enabled: {default_quant_conv.is_output_quant_enabled}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we now try to pass in a random floating-point tensor as input, as expected we get the output of the convolution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.2594,  0.5392,  0.5916],\n",
       "          [ 0.3493,  0.6813,  0.2499],\n",
       "          [ 1.3732,  0.1229, -0.0084]],\n",
       "\n",
       "         [[ 0.0031, -0.1702,  0.1069],\n",
       "          [-0.8181, -0.8056,  0.0385],\n",
       "          [-0.4738,  0.0589,  0.1278]],\n",
       "\n",
       "         [[-0.1718, -0.1162, -0.1526],\n",
       "          [-0.9903, -0.3541,  0.1645],\n",
       "          [ 0.0557, -0.4458, -0.2080]]]], grad_fn=<ThnnConv2DBackward0>)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "out = default_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case we are computing the convolution between an unquantized input tensor and quantized weights, so the output in general is unquantized."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A QuantConv2d with quantization disabled everywhere behaves like a standard `Conv2d`. Again can easily verify this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Conv2d\n",
    "\n",
    "torch.manual_seed(0)  # set a seed to make sure the random weight init is reproducible\n",
    "disabled_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, weight_quant=None)\n",
    "torch.manual_seed(0)  # reproduce the same random weight init as above\n",
    "float_conv = Conv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False)\n",
    "inp = torch.randn(1, 2, 5, 5)\n",
    "assert torch.isclose(disabled_quant_conv(inp), float_conv(inp)).all().item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we have just seen, Brevitas allows users as much freedom as possible to experiment with quantization, meaning that computation between quantized and unquantized values is considered legal. This allows users to mix Brevitas layers with Pytorch layers with little restrictions.  \n",
    "To make this possible, quantized values are typically represented in *dequantized format*, meaning that - in the case of affine quantization implemented in Brevitas - zero-point and scale factor are applied to their integer values according to the formula **quant_value = (integer_value - zero_point) * scale**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QuantTensor\n",
    "\n",
    "We can directly observe the quantized weights by calling the weight quantizer on the layer's weights: `default_quant_conv.weight_quant(quant_conv.weight)`, which for shortness is already implemented as `default_quant_conv.quant_weight()` :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[-0.0790,  0.0503, -0.0934],\n",
       "          [-0.1149, -0.1903, -0.1329],\n",
       "          [-0.1813,  0.0108,  0.0593]],\n",
       "\n",
       "         [[ 0.0970, -0.0215, -0.0144],\n",
       "          [ 0.2280,  0.1239, -0.0090],\n",
       "          [ 0.1957, -0.2011, -0.0108]]],\n",
       "\n",
       "\n",
       "        [[[-0.0018, -0.1957,  0.1993],\n",
       "          [-0.0359,  0.1778, -0.1400],\n",
       "          [ 0.0916,  0.1059,  0.2173]],\n",
       "\n",
       "         [[-0.1670,  0.1939, -0.2191],\n",
       "          [-0.0215,  0.1688, -0.1383],\n",
       "          [-0.0449, -0.1185,  0.1742]]],\n",
       "\n",
       "\n",
       "        [[[-0.0808, -0.1652, -0.0233],\n",
       "          [-0.0700,  0.0467, -0.0485],\n",
       "          [ 0.1059,  0.1418,  0.1077]],\n",
       "\n",
       "         [[-0.0593,  0.0108,  0.0036],\n",
       "          [-0.1508,  0.0808,  0.1616],\n",
       "          [ 0.0144, -0.0287, -0.1365]]]], grad_fn=<MulBackward0>), scale=tensor(0.0018, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "default_quant_conv.quant_weight()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the quantized weights are wrapped in a data structure implemented by Brevitas called `QuantTensor`. A `QuantTensor` is a way to represent an affine quantized tensor with all its metadata, meaning: the `value` of the quantized tensor in *dequantized* format, `scale`, `zero_point`, `bit_width`, whether the quantized value it's `signed` or not, and whether the tensor was generated in `training` mode. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, we have that the quantized value (in dequantized format) can be computer from its integer representation, together with zero-point and scale:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "int_weight = default_quant_conv.int_weight()\n",
    "zero_point = default_quant_conv.quant_weight_zero_point()\n",
    "scale = default_quant_conv.quant_weight_scale()\n",
    "quant_weight_manually = (int_weight - zero_point) * scale\n",
    "\n",
    "assert default_quant_conv.quant_weight().value.isclose(quant_weight_manually).all().item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A *valid* QuantTensor correctly populates all its fields with values `!= None` and respect the **affine quantization invariant**, i.e. `value / scale + zero_point` is (accounting for rounding errors) an *integer* that can be represented within the interval defined by the `bit_width` and `signed` fields of the `QuantTensor`. A *non-valid* one doesn't.\n",
    "We can observe that the quantized weights are indeed marked as valid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert default_quant_conv.quant_weight().is_valid"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calling `is_valid` is relative expensive, so it should be using sparingly, but there are a few cases where a non-valid QuantTensor might be generated that is important to be aware of. Say we have two QuantTensor as output of the same quantized activation, and we want to sum them together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0173, grad_fn=<DivBackward0>)\n",
      "tensor(0.0307, grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "from brevitas.quant_tensor import QuantTensor\n",
    "\n",
    "quant_act = QuantIdentity(return_quant_tensor=True)\n",
    "\n",
    "out_tensor_0 = quant_act(torch.randn(1,2,5,5))\n",
    "out_tensor_1 = quant_act(torch.randn(1,2,5,5))\n",
    "\n",
    "assert out_tensor_0.is_valid\n",
    "assert out_tensor_1.is_valid\n",
    "print(out_tensor_0.scale)\n",
    "print(out_tensor_1.scale)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both QuantTensor are valid but since the quantized activation is in training mode by default, their scale factors are going to be different. It is important to note that the behaviour is different at evaluation time, where the two scale factors will be the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 0.9489, -0.9111, -0.0536,  0.5788,  0.3645],\n",
       "          [ 0.3401,  1.4325,  0.6498,  0.6411, -1.4390],\n",
       "          [-1.9029,  0.7012,  0.1591,  1.9235,  0.5883],\n",
       "          [-2.7258,  2.5330,  0.9165, -0.0820,  3.4148],\n",
       "          [-0.3651,  1.0164,  0.9567, -0.2758, -1.1376]],\n",
       "\n",
       "         [[-0.2414,  2.2111, -1.9124, -2.3814, -0.8805],\n",
       "          [ 1.3191, -0.8965, -0.2048, -3.8113,  1.1142],\n",
       "          [-0.3381, -0.2238,  1.2661,  0.0068,  0.2567],\n",
       "          [ 0.0731, -0.4280,  0.0909,  0.0875, -1.6851],\n",
       "          [-0.7744, -1.4127, -0.8143,  1.3557, -0.2802]]]],\n",
       "       grad_fn=<AddBackward0>), scale=tensor(0.0240, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(9.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_tensor = out_tensor_0 + out_tensor_1\n",
    "out_tensor"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because we set `training` to `True` for both of them, we are allowed to sum them even if they have different scale factors. The output QuantTensor will have the correct `bit_width`, and a scale which is the average of the two original scale factors. This is done only at training time, in order to propagate gradient information, however the consequence is that the resulting QuantTensor is no longer valid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not out_tensor.is_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QuantTensor` implements `__torch_function__` to handle being called from torch functional operators (e.g. ops under `torch.nn.functional`). Passing a QuantTensor to supported ops that are invariant to quantization, e.g. max-pooling, preserve the the validity of a QuantTensor. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[1.5800, 1.0157],\n",
       "          [1.4445, 0.8577]],\n",
       "\n",
       "         [[0.5643, 1.2414],\n",
       "          [1.0383, 0.9028]],\n",
       "\n",
       "         [[0.5191, 0.6546],\n",
       "          [2.1442, 0.5868]]]], grad_fn=<MaxPool2DWithIndicesBackward0>), scale=tensor(0.0226, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "quant_identity = QuantIdentity(return_quant_tensor=True)\n",
    "quant_tensor = quant_identity(torch.randn(1, 3, 4, 4))\n",
    "torch.nn.functional.max_pool2d(quant_tensor, kernel_size=2, stride=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For ops that are not invariant to quantization, a `QuantTensor` decays into a floating-point `torch.Tensor`. Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.4943, -0.9938, -0.9073,  0.7681],\n",
       "          [-0.3262,  0.9186,  0.1786,  0.3659],\n",
       "          [ 0.7489,  0.8946, -0.0451, -0.5594],\n",
       "          [-0.1346, -0.4943, -0.4770,  0.6951]],\n",
       "\n",
       "         [[ 0.0676,  0.5111,  0.4943,  0.8459],\n",
       "          [-0.8990, -0.9426,  0.0676, -0.7945],\n",
       "          [-0.9220,  0.0676, -0.5594,  0.6321],\n",
       "          [-0.0676,  0.7772,  0.7177, -0.4414]],\n",
       "\n",
       "         [[ 0.4770,  0.2220,  0.0676,  0.5747],\n",
       "          [-0.0451, -0.6710, -0.4594, -0.3462],\n",
       "          [ 0.9729, -0.7177, -0.5896, -0.5276],\n",
       "          [-0.0900,  0.8852,  0.5276, -0.4414]]]], grad_fn=<TanhBackward0>)"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.tanh(quant_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input Quantization\n",
    "\n",
    "We can obtain a valid output `QuantTensor` by making sure that both input and weight of `QuantConv2d` are quantized. To do so, we can set a quantizer for `input_quant`. In this example we pick a *signed 8-bit* quantizer with *per-tensor floating-point scale factor*:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 0.9693, -0.9431,  0.2459],\n",
       "          [ 0.5416,  0.9037, -0.5278],\n",
       "          [-0.6207, -1.3578, -0.4815]],\n",
       "\n",
       "         [[ 0.4551, -1.4065,  0.8889],\n",
       "          [-0.3393,  0.0803, -0.1748],\n",
       "          [-0.0977,  0.6284, -0.7193]],\n",
       "\n",
       "         [[ 0.3655,  0.7626, -0.2634],\n",
       "          [-0.3453,  0.3349,  0.1923],\n",
       "          [ 0.5993, -0.9579,  0.3557]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[3.2208e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8ActPerTensorFloat\n",
    "\n",
    "input_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, \n",
    "    input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)\n",
    "out_tensor = input_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert out_tensor.is_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens internally is that the input tensor passed to `input_quant_conv` is being quantized before being passed to the convolution operator. That means we are now computing a convolution between two quantized tensors, which mimplies that the output of the operation is also quantized. As expected then `out_tensor` is marked as valid. \n",
    "\n",
    "Another important thing to notice is how the `bit_width` field of `out_tensor` is relatively high at *21 bits*. In Brevitas, the assumption is always that the output bit-width of an operator reflects the worst-case size of the *accumulator* required by that operation. In other terms, given the *size* of the input and weight tensors and their *bit-widths*, 21 is the bit-width that would be required to represent the largest possible output value that could be generated. This makes sure that the affine quantization invariant is always respected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could have obtained a similar result by directly passing as input a QuantTensor. In this example we are directly defining a QuantTensor ourselves, but it could also be the output of a previous layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 5.7000e-03,  2.5000e-03, -1.2400e-02, -7.2000e-03,  3.7000e-03],\n",
       "          [-2.3000e-03,  7.0000e-04, -1.2700e-02,  5.2000e-03,  4.0000e-04],\n",
       "          [-7.9000e-03,  9.5000e-03,  6.6000e-03,  5.4000e-03,  2.5000e-03],\n",
       "          [ 1.1100e-02,  2.4000e-03,  1.0000e-02, -3.7000e-03,  7.2000e-03],\n",
       "          [-1.1500e-02, -5.8000e-03, -9.3000e-03,  1.0000e-02,  3.5000e-03]],\n",
       "\n",
       "         [[-6.8000e-03,  1.1500e-02, -1.0600e-02, -1.5000e-03, -1.9000e-03],\n",
       "          [ 2.9000e-03,  9.5000e-03,  7.2000e-03, -3.7000e-03,  7.7000e-03],\n",
       "          [-2.4000e-03, -8.9000e-03, -1.2000e-02, -8.1000e-03,  7.2000e-03],\n",
       "          [-1.1300e-02, -9.7000e-03, -1.0000e-03,  1.0100e-02,  3.8000e-03],\n",
       "          [-1.1900e-02,  6.9000e-03,  8.3000e-03,  1.0000e-04, -6.9000e-03]]]]), scale=tensor(1.0000e-04), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant_tensor import QuantTensor\n",
    "\n",
    "scale = 0.0001\n",
    "bit_width = 8\n",
    "zero_point = 0.\n",
    "int_value = torch.randint(low=- 2 ** (bit_width - 1), high=2 ** (bit_width - 1) - 1, size=(1, 2, 5, 5))\n",
    "quant_value = (int_value - zero_point) * scale\n",
    "quant_tensor_input = QuantTensor(\n",
    "    quant_value, \n",
    "    scale=torch.tensor(scale), \n",
    "    zero_point=torch.tensor(zero_point), \n",
    "    bit_width=torch.tensor(float(bit_width)),\n",
    "    signed=True,\n",
    "    training=True)\n",
    "quant_tensor_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert quant_tensor_input.is_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: how we are explicitly forcing `value`, `scale`, `zero_point` and `bit_width` to be floating-point `torch.Tensor`, as this is expected by Brevitas but it's currently not enforced automatically at initialization time.\n",
    "\n",
    "If we now pass in `quant_tensor_input` to `return_quant_conv`, we will see that indeed the output is a valid 21-bit `QuantTensor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 0.0085,  0.0066,  0.0050],\n",
       "          [-0.0038, -0.0009, -0.0115],\n",
       "          [-0.0055, -0.0037,  0.0009]],\n",
       "\n",
       "         [[ 0.0015, -0.0027, -0.0079],\n",
       "          [-0.0034, -0.0060,  0.0043],\n",
       "          [-0.0008,  0.0052, -0.0033]],\n",
       "\n",
       "         [[-0.0015,  0.0082, -0.0038],\n",
       "          [-0.0021,  0.0004, -0.0054],\n",
       "          [-0.0021, -0.0079,  0.0013]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.8448e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "return_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, return_quant_tensor=True)\n",
    "out_tensor = return_quant_conv(quant_tensor_input)\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert out_tensor.is_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also pass in an input `QuantTensor` to a layer that has `input_quant` enabled. In that case, the input gets re-quantized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[-0.0035, -0.0037, -0.0050],\n",
       "          [ 0.0010, -0.0051, -0.0027],\n",
       "          [-0.0010,  0.0047,  0.0017]],\n",
       "\n",
       "         [[ 0.0021,  0.0002,  0.0027],\n",
       "          [ 0.0028,  0.0002, -0.0044],\n",
       "          [ 0.0008, -0.0052, -0.0024]],\n",
       "\n",
       "         [[ 0.0010, -0.0052, -0.0011],\n",
       "          [-0.0018,  0.0024,  0.0011],\n",
       "          [-0.0001,  0.0039,  0.0035]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.7410e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_quant_conv(quant_tensor_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output Quantization\n",
    "\n",
    "Let's now look at would have happened if we instead enabled output quantization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 0.2111,  0.4060,  0.3654],\n",
       "          [-0.7876,  0.8119, -0.9825],\n",
       "          [-0.5115,  0.3979, -0.3248]],\n",
       "\n",
       "         [[ 0.3816,  0.0568, -0.0812],\n",
       "          [ 1.0312, -0.7876,  0.8038],\n",
       "          [-0.3491, -0.4141,  0.0650]],\n",
       "\n",
       "         [[-0.5846, -0.4222, -0.0731],\n",
       "          [-0.7389,  0.5034, -0.2517],\n",
       "          [-0.1624, -0.4385,  0.7308]]]], grad_fn=<MulBackward0>), scale=tensor(0.0081, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8ActPerTensorFloat\n",
    "\n",
    "output_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, \n",
    "    output_quant=Int8ActPerTensorFloat, return_quant_tensor=True)\n",
    "out_tensor = output_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert out_tensor.is_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see again that the output is a valid `QuantTensor`. However, what happened internally is quite different from before.  \n",
    "Previously, we computed the convolution between two quantized tensors, and got a quantized tensor as output.  \n",
    "In this case instead, we compute the convolution between a quantized and an unquantized tensor, we take its unquantized output and we quantize it.  \n",
    "The difference is obvious once we look at the output `bit_width`. In the previous case, we had that the `bit_width` reflected the size of the output accumulator. In this case instead, we have `bit_width=tensor(8.)`, which is what we expected since `output_quant` had been set to an *Int8* quantizer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bias Quantization\n",
    "\n",
    "There is an important scenario where the various options we just saw make a practical difference, and it's quantization of *bias*. In many contexts, such as in the ONNX standard opset and in FINN, bias is assumed to be quantized with scale factor equal to *input scale * weight scale*, which means that we need a valid quantized input somehow. A predefined bias quantizer that reflects that assumption is `brevitas.quant.scaled_int.Int8Bias`. If we simply tried to set it to a `QuantConv2d` without any sort of input quantization, we would get an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Input scale required",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_48365/2280634207.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0min_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     bias_quant=Int8Bias, return_quant_tensor=True)\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mbias_quant_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/opt/conda/envs/torch_1.10/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1100\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1101\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1102\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1103\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1104\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/workspace/scratch/git/fork_brevitas/src/brevitas/nn/quant_conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mQuantTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mQuantTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 192\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    193\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    194\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0minner_forward_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquant_weight\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquant_bias\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/workspace/scratch/git/fork_brevitas/src/brevitas/nn/quant_layer.py\u001b[0m in \u001b[0;36mforward_impl\u001b[0;34m(self, inp)\u001b[0m\n\u001b[1;32m    330\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    331\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 332\u001b[0;31m             \u001b[0mquant_bias\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias_quant\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_scale\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_bit_width\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    333\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcache_inference_quant_bias\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    334\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_cached_bias\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_CachedIO\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquant_bias\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmetadata_only\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/torch_1.10/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1100\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1101\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1102\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1103\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1104\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/workspace/scratch/git/fork_brevitas/src/brevitas/proxy/parameter_quant.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, input_scale, input_bit_width)\u001b[0m\n\u001b[1;32m    160\u001b[0m             \u001b[0mimpl\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexport_handler\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexport_mode\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor_quant\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    161\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_input_scale\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0minput_scale\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 162\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input scale required\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    163\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_input_bit_width\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0minput_bit_width\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    164\u001b[0m                 \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input bit-width required\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Input scale required"
     ]
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8Bias\n",
    "\n",
    "bias_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    bias_quant=Int8Bias, return_quant_tensor=True)\n",
    "bias_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can solve the issue by passing in a valid `QuantTensor`, e.g. the `quant_tensor_input`  we defined above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 0.0005,  0.0043, -0.0004],\n",
       "          [ 0.0005,  0.0106,  0.0012],\n",
       "          [ 0.0021,  0.0007, -0.0050]],\n",
       "\n",
       "         [[-0.0067, -0.0035, -0.0059],\n",
       "          [-0.0050, -0.0015, -0.0039],\n",
       "          [ 0.0015,  0.0028, -0.0008]],\n",
       "\n",
       "         [[-0.0051, -0.0050,  0.0060],\n",
       "          [-0.0015,  0.0037,  0.0071],\n",
       "          [ 0.0067,  0.0035, -0.0071]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.8108e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 120,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bias_quant_conv(quant_tensor_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or by enabling input quantization and then passing in a float a `torch.Tensor` or a `QuantTensor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[-0.3825,  0.1371,  0.9135],\n",
       "          [-0.2016,  0.7495, -0.4071],\n",
       "          [-0.0755,  0.5283,  0.2388]],\n",
       "\n",
       "         [[ 0.0788, -0.3802, -0.2234],\n",
       "          [ 0.8678, -0.5546,  0.4408],\n",
       "          [-0.6788,  0.4422,  0.3007]],\n",
       "\n",
       "         [[ 0.4412, -0.3205,  1.0033],\n",
       "          [-0.0083, -0.3295, -0.2076],\n",
       "          [ 0.4417, -0.1046, -0.3493]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[3.8610e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_bias_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    input_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)\n",
    "input_bias_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[ 0.0036,  0.0024, -0.0033],\n",
       "          [ 0.0050,  0.0080, -0.0014],\n",
       "          [-0.0036, -0.0080, -0.0029]],\n",
       "\n",
       "         [[ 0.0083, -0.0093,  0.0048],\n",
       "          [ 0.0035,  0.0015, -0.0011],\n",
       "          [-0.0003,  0.0067,  0.0013]],\n",
       "\n",
       "         [[-0.0009, -0.0019,  0.0039],\n",
       "          [ 0.0010,  0.0056, -0.0037],\n",
       "          [ 0.0091, -0.0095,  0.0054]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.8384e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_bias_quant_conv(quant_tensor_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the output `bit_width=tensor(22.)`. This is because, in the worst-case, summing a *21-bit* integer (the size of the accumulator before bias is added) and an *8-bit* integer (the size of quantized bias) gives a *22-bit* integer.\n",
    "\n",
    "Let's try now to enable output quantization instead of input quantization. That wouldn't have solved the problem with bias quantization, as output quantization is performed after bias is added:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Input scale required",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_48365/2990591641.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0min_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_channels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkernel_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m     output_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0moutput_bias_quant_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/opt/conda/envs/torch_1.10/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1100\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1101\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1102\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1103\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1104\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/workspace/scratch/git/fork_brevitas/src/brevitas/nn/quant_conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mQuantTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mQuantTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 192\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    193\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    194\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0minner_forward_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquant_weight\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquant_bias\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/workspace/scratch/git/fork_brevitas/src/brevitas/nn/quant_layer.py\u001b[0m in \u001b[0;36mforward_impl\u001b[0;34m(self, inp)\u001b[0m\n\u001b[1;32m    330\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    331\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 332\u001b[0;31m             \u001b[0mquant_bias\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias_quant\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_scale\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_bit_width\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    333\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcache_inference_quant_bias\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    334\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_cached_bias\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_CachedIO\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquant_bias\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmetadata_only\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/torch_1.10/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1100\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1101\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1102\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1103\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1104\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/workspace/scratch/git/fork_brevitas/src/brevitas/proxy/parameter_quant.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, input_scale, input_bit_width)\u001b[0m\n\u001b[1;32m    160\u001b[0m             \u001b[0mimpl\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexport_handler\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexport_mode\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor_quant\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    161\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_input_scale\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0minput_scale\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 162\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input scale required\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    163\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_input_bit_width\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0minput_bit_width\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    164\u001b[0m                 \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Input bit-width required\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Input scale required"
     ]
    }
   ],
   "source": [
    "output_bias_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    output_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)\n",
    "output_bias_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not all scenarios require bias quantization to depend on the scale factor of the input. In those cases, biases can be quantized the same way weights are quantized, and have their own scale factor. In Brevitas, a predefined quantizer that reflects this other scenario is `Int8BiasPerTensorFloatInternalScaling`. In this case then a valid quantized input is not required:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.2152,  0.8346,  0.0746],\n",
       "          [-0.0738, -0.5212,  0.1019],\n",
       "          [-0.6004,  0.1500, -0.1453]],\n",
       "\n",
       "         [[-1.1551, -1.3458, -0.1312],\n",
       "          [ 0.2502, -0.5267,  0.2412],\n",
       "          [-0.3556, -0.3289, -0.2276]],\n",
       "\n",
       "         [[-0.4599, -0.6094,  0.4682],\n",
       "          [-0.5064, -0.6768, -0.6638],\n",
       "          [ 0.0066, -0.3581,  0.2359]]]], grad_fn=<ThnnConv2DBackward0>)"
      ]
     },
     "execution_count": 124,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from brevitas.quant.scaled_int import Int8BiasPerTensorFloatInternalScaling\n",
    "\n",
    "bias_internal_scale_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    bias_quant=Int8BiasPerTensorFloatInternalScaling, return_quant_tensor=False)\n",
    "bias_internal_scale_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are a couple of situations to be aware of concerning bias quantization that can lead to changes in the output `zero_point`.\n",
    "\n",
    "Let's consider the scenario where we compute the convolution between a quantized input tensor and quantized weights. In the first case, we then add an *unquantized* bias on top of the output. In the second one, we add a bias quantized with its own scale factor, e.g. with the `Int8BiasPerTensorFloatInternalScaling` quantizer. In both cases, in order to make sure the output `QuantTensor` is valid (i.e. the affine quantization invariant is respected), the output `zero_point` becomes non-zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantTensor(value=tensor([[[[-0.6879, -0.6632, -0.2411],\n",
       "          [ 0.2064, -0.7371,  0.3910],\n",
       "          [ 0.9533,  0.2994,  0.6546]],\n",
       "\n",
       "         [[-0.4684, -0.4495, -0.5021],\n",
       "          [ 0.5738,  0.4199, -0.3380],\n",
       "          [ 0.6218, -0.0408, -0.8483]],\n",
       "\n",
       "         [[-0.5625,  0.1837, -1.0575],\n",
       "          [-1.2816, -0.4993, -0.3409],\n",
       "          [ 0.4556, -1.4269,  0.5369]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[3.0975e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor([[[[ 1276.0774]],\n",
       "\n",
       "         [[-3152.4585]],\n",
       "\n",
       "         [[ 7320.2324]]]], grad_fn=<DivBackward0>), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))"
      ]
     },
     "execution_count": 125,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unquant_bias_input_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)\n",
    "out_tensor = unquant_bias_input_quant_conv(torch.randn(1, 2, 5, 5))\n",
    "out_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert out_tensor.is_valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, an important point about `QuantTensor`. With the exception of learned bit-width (which will be the subject of a separate tutorial) and some of the bias quantization scenarios we have just seen, usually returing a `QuantTensor` is not necessary and can create extra complexity. This is why currently `return_quant_tensor` defaults to `False`. We can easily see it in an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ 0.8357,  0.0733,  0.9527],\n",
       "          [ 0.1803,  0.2154,  0.7598],\n",
       "          [ 1.1121, -0.8728,  1.0039]],\n",
       "\n",
       "         [[ 0.7917,  1.0063,  0.6516],\n",
       "          [-0.1852, -0.7263,  0.0956],\n",
       "          [-0.1876,  0.2747, -0.1617]],\n",
       "\n",
       "         [[ 0.8299,  0.9934, -0.3821],\n",
       "          [ 0.4865,  0.9309, -0.7924],\n",
       "          [-0.4201,  0.2343,  0.1532]]]], grad_fn=<ThnnConv2DBackward0>)"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bias_input_quant_conv = QuantConv2d(\n",
    "    in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,\n",
    "    input_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias)\n",
    "bias_input_quant_conv(torch.randn(1, 2, 5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Altough not obvious, the output is actually implicitly quantized."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_1.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "98d39f4fefefd06bad749c114549705efc94e31e1348b8fb7b25329049b354fa"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
